MODULE fefmm_box_class

   USE fmm_global_paras
   IMPLICIT NONE
   PRIVATE
   ! Public procedures
   PUBLIC :: fefmm_init_box_class,           &
             fefmm_free_box_class,           &
             fefmm_update_J_matrix

   INTEGER(INTK) :: FEdim, FEdimB, NFS
   INTEGER(INTK) :: LIPN
   REAL(REALK)   :: FEgrain
   REAL(REALK)   :: box_centre(3) 
   REAL(REALK)   :: box_grain

   !REAL(REALK), ALLOCATABLE :: FEdens(:)
   REAL(REALK), ALLOCATABLE :: FEop(:,:)
   REAL(REALK), ALLOCATABLE :: FEpot(:)
   REAL(REALK), POINTER :: FEdens(:)

CONTAINS

!-------------------------------------------------------------------------------

   SUBROUTINE fefmm_init_box_class(scheme,basis,sh_pairs,dens)

      USE fefmm_density,  ONLY: fefmm_init_density
      USE fefmm_operator, ONLY: fefmm_get_FEop

      IMPLICIT NONE
      TYPE(fmm_basis),    INTENT(IN) :: basis
      TYPE(scheme_paras), INTENT(IN) :: scheme
      TYPE(fmm_sh_pairs), INTENT(IN) :: sh_pairs(:)
      REAL(REALK),        INTENT(IN) :: dens(:,:)

      FEdim   = scheme%FEdim   ! number of grid points in one box (1D)
!      FEdimB  = (3*FEdim -2)   ! number grid points in near-field space (1D)
!      NFS     = 3              ! number boxes in near-field space (1D)
      FEdimB  = FEdim
      NFS     = 1

      LIPN = scheme%lipn
      FEgrain = scheme%grain / (FEdim-1)
      box_grain = NFS*scheme%grain
      box_centre(:) = zero

      CALL fefmm_init_density(scheme,basis,sh_pairs,dens)
      NULLIFY(FEdens)
!      ALLOCATE (FEdens(FEdimB**3))
      ALLOCATE (FEpot(FEdim**3))

!print *, "Getting FEop..."
!      ALLOCATE (FEop(n,FEdim**3))
!      PRINT '(A,F10.2,A)', " Memory for FEop = ", SIZE(FEOp)*8d0*1d-6, " MB"
!      CALL fefmm_get_FEop(FEop)

   END SUBROUTINE fefmm_init_box_class

!-------------------------------------------------------------------------------

   SUBROUTINE fefmm_free_box_class

      IMPLICIT NONE

!      IF (ALLOCATED (FEdens)) DEALLOCATE (FEdens)
      IF (ALLOCATED (FEop)) DEALLOCATE (FEop)
      IF (ALLOCATED (FEpot)) DEALLOCATE (FEpot)

   END SUBROUTINE fefmm_free_box_class

!-------------------------------------------------------------------------------
!
!   SUBROUTINE fefmm_build_potential(centre,nel)
!
!      USE fefmm_density,   ONLY: fefmm_get_FEdens
!      USE fefmm_LIP_trace, ONLY: fefmm_trace
!
!      IMPLICIT NONE
!      REAL(REALK), INTENT(IN)  :: centre(3)
!      REAL(REALK), INTENT(OUT) :: nel
!
!!REAL(REALK) :: R
!
!      box_centre(:) = centre(:)
!
!print *, "Building FEdens..."
!      CALL fefmm_get_FEdens(box_centre,FEdens)
!!      CALL fefmm_model_FEdens
!!      FEdens = one
!      CALL fefmm_trace(FEdens,FEdim,LIPN,FEgrain,nel)
!!      IF (ABS(nel) < 1d-15) RETURN    ! broken???
!!print *, "# electrons in box = ", nel
!
!print *, "Building FEpot..."
!!      FEpot(:) = one 
!!      FEpot(:) = MATMUL(FEdens(:),FEop(:,:))
!      CALL fefmm_build_FEpot_direct
!
!!if (abs(fepot(1)) > 1d-5) then
!!R = sqrt(dot_product(centre,centre))
!!print *, 'Rpq:', R
!!print *, 'FEpot =', FEpot(1)
!!print *, 'exact =', erf(sqrt(2d0)*R)/R
!!print *
!!end if
!!stop
!
!   END SUBROUTINE fefmm_build_potential
!
!-------------------------------------------------------------------------------

   SUBROUTINE fefmm_get_Fij(t,Fij)

      USE fefmm_LIP1_Fij, ONLY: fefmm_build_Fij_LIP1
      USE fefmm_LIP2_Fij, ONLY: fefmm_build_Fij_LIP2
      USE fefmm_LIP3_Fij, ONLY: fefmm_build_Fij_LIP3
      USE fefmm_LIP6_Fij, ONLY: fefmm_build_Fij_LIP6

      IMPLICIT NONE
      REAL(REALK), INTENT(IN)  :: t
      REAL(REALK), INTENT(OUT) :: Fij(FEdim,FEdim)

      SELECT CASE (LIPN)
         CASE (1)
            CALL fefmm_build_Fij_LIP1(FEdim,NFS,FEgrain,t,Fij)
         CASE (2)
            CALL fefmm_build_Fij_LIP2(FEdim,NFS,FEgrain,t,Fij)
         CASE (3)
            CALL fefmm_build_Fij_LIP3(FEdim,NFS,FEgrain,t,Fij)
         CASE (6)
            CALL fefmm_build_Fij_LIP6(FEdim,NFS,FEgrain,t,Fij)
         CASE DEFAULT
            CALL fmm_quit('NLIP too large in get_Fij!')
      END SELECT

   END SUBROUTINE fefmm_get_Fij

!-------------------------------------------------------------------------------

   SUBROUTINE fefmm_build_FEpot_direct

      USE fefmm_quadrature, ONLY: fefmm_get_roots_and_weights

      IMPLICIT NONE
      INTEGER(INTK), PARAMETER :: npoints = 500
      REAL(REALK),   PARAMETER :: tmax = 10000.0d0

      REAL(REALK)   :: t(npoints), wt(npoints)
      REAL(REALK)   :: dens3D(FEdim,FEdim,FEdim)
      REAL(REALK)   :: Fij(FEdim,FEdim)
      REAL(REALK)   :: V(FEdim,FEdim,FEdim)
      REAL(REALK)   :: W(FEdim,FEdim,FEdim)
      INTEGER(INTK) :: nt, ix,iy,iz

      dens3D = RESHAPE(FEdens,(/FEdim,FEdim,FEdim/)) ! -> FEdens(z,y,x)

      CALL fefmm_get_roots_and_weights(npoints,tmax,t,wt)
      FEpot(:) = zero

      DO nt = 1, npoints
!print *, nt, FEgrain, t(nt), wt(nt)

         CALL fefmm_get_Fij(t(nt),Fij)

         DO ix = 1, FEdim
            DO iy = 1, FEdim
               !V(:,iy,ix) = MATMUL(dens3D(:,iy,ix),Fij(:,:))
               V(:,iy,ix) = MATMUL(dens3D(:,iy,ix),Fij)
            END DO
         END DO

         DO ix = 1, FEdim
            DO iz = 1, FEdim
               W(iz,:,ix) = MATMUL(V(iz,:,ix),Fij)
            END DO
         END DO

         DO iy = 1, FEdim
            DO iz = 1, FEdim
               V(iz,iy,:) = MATMUL(W(iz,iy,:),Fij)
            END DO
         END DO

         FEpot(:) = FEpot(:) + wt(nt)*RESHAPE(V,(/FEdim**3/))

      END DO

      FEpot(:) = FEpot(:)*two/ROOTPI

   END SUBROUTINE fefmm_build_FEpot_direct

!-------------------------------------------------------------------------------
!
!   SUBROUTINE fefmm_contract_FE_potential(scheme,basis,sh_pairs,box,J_matrix)
!
!      USE fefmm_FE_ints, ONLY: fefmm_init_FE_ints,             &
!                               fefmm_free_FE_ints,             &
!                               fefmm_contract_FEpot
!
!      IMPLICIT NONE
!      TYPE(scheme_paras), INTENT(IN)    :: scheme
!      TYPE(fmm_basis),    INTENT(IN)    :: basis
!      TYPE(fmm_sh_pairs), INTENT(IN)    :: sh_pairs(:)
!      REAL(REALK),        INTENT(IN)    :: box(3)
!      REAL(REALK),        INTENT(INOUT) :: J_matrix(:,:)
!
!print *, "Contracting FEpot..."
!      CALL fefmm_init_FE_ints(basis,FEdim,LIPN)
!      CALL fefmm_contract_FEpot(scheme,basis,sh_pairs,box,FEpot,J_matrix)
!      CALL fefmm_free_FE_ints
!
!   END SUBROUTINE fefmm_contract_FE_potential
!
!-------------------------------------------------------------------------------

   SUBROUTINE fefmm_update_J_matrix(box,nel,J_matrix,nbas)

      USE fefmm_density,   ONLY: fefmm_get_projAO, fefmm_get_FEdens
      USE fefmm_LIP_inner, ONLY: fefmm_init_inner, fefmm_free_inner,   &
                                 fefmm_inner, fefmm_trace

      IMPLICIT NONE
      REAL(REALK), INTENT(IN)    :: box(3)
      REAL(REALK), INTENT(OUT)   :: nel
      REAL(REALK), INTENT(INOUT) :: J_matrix(nbas,nbas)
      INTEGER(INTK), INTENT(IN)  :: nbas

      REAL(REALK), POINTER :: projAO_p(:)
      REAL(REALK), POINTER :: projAO_q(:)
      REAL(REALK)   :: projAO_pq(FEdim**3)
      REAL(REALK)   :: tmp
      INTEGER(INTK) :: ijk, p,q

      NULLIFY(projAO_p, projAO_q)

print *, "Building FEdens..."
      CALL fefmm_get_FEdens(box,FEdens)
      CALL fefmm_init_inner(FEdim,LIPN,FEgrain)
      CALL fefmm_trace(FEdens,nel)
!      IF (ABS(nel) < 1d-15) RETURN    ! broken???
!print *, "# electrons in box = ", nel

print *, "Building FEpot..."
!      FEpot(:) = one 
!      FEpot(:) = MATMUL(FEdens(:),FEop(:,:))
      CALL fefmm_build_FEpot_direct

print *, "Contracting FEpot..."
      DO p = 1, nbas
         CALL fefmm_get_projAO(p,projAO_p)

         DO q = 1, nbas
            CALL fefmm_get_projAO(q,projAO_q)

            DO ijk = 1, FEdim**3
               projAO_pq(ijk) = projAO_p(ijk) * projAO_q(ijk)
            END DO

            CALL fefmm_inner(projAO_pq,FEpot,tmp)
            J_matrix(p,q) = J_matrix(p,q) + tmp

         END DO
      END DO

      NULLIFY(projAO_p, projAO_q)
      CALL fefmm_free_inner

   END SUBROUTINE fefmm_update_J_matrix

!-------------------------------------------------------------------------------

END MODULE fefmm_box_class

